"""
@Filename       : twitter_preprocess.py
@Create Time    : 2021/3/16 14:51
@Author         : Rylynn
@Description    : 

"""
import os
import random

import pandas as pd


def load_graph(rootpath, file, keep_user):
    filepath = os.path.join(rootpath, file)
    graph_file = open(os.path.join(rootpath, 'processed', 'graph.txt'), mode='w+', encoding='utf8')

    with open(filepath, 'r+') as f:
        for line in f.readlines():
            mutual, time, source, target = line.split(',')
            mutual = eval(mutual)
            source = eval(source)
            target = eval(target)
            if source not in keep_user or target not in keep_user:
                continue
            if mutual == 1:
                graph_file.write('{} {}\n'.format(source, target))
                graph_file.write('{} {}\n'.format(target, source))
            else:
                graph_file.write('{} {}\n'.format(source, target))


def load_diffusion_data(rootpath, file):
    filepath = os.path.join(rootpath, file)
    user_freq_dict = {}
    with open(filepath, 'r') as f:
        for line in f.readlines():
            time, user, story = line.split(',')
            time = eval(time)
            user = eval(user)
            story = eval(story)
            if not user_freq_dict.get(user):
                user_freq_dict[user] = 1
            else:
                user_freq_dict[user] = user_freq_dict[user] + 1

    filter_threshold = 55
    filtered_user_dict = filter(lambda kv: kv[1] >= filter_threshold, user_freq_dict.items())
    keep_users = set(dict(filtered_user_dict).keys())

    story_users_dict = dict()
    with open(filepath, 'r') as f:
        for line in f.readlines():
            time, user, story = line.split(',')
            time = eval(time)
            user = eval(user)
            story = eval(story)
            if not story_users_dict.get(story):
                if user in keep_users:
                    story_users_dict[story] = [(user, time)]
            else:
                story_users_dict[story].append((user, time))

    story_list = list(story_users_dict.keys())
    random.shuffle(story_list)

    training_file = open(os.path.join(rootpath, 'processed', 'train.txt'), 'w+', encoding='utf8')
    testing_file = open(os.path.join(rootpath, 'processed', 'test.txt'), 'w+', encoding='utf8')
    for idx, story in enumerate(story_list):
        cascade = story_users_dict[story]
        if idx < len(story_list) * 0.75:
            training_file.write('{}'.format(cascade[0][0]))
            for user, time in cascade[1:]:
                if user not in keep_users:
                    continue
                training_file.write(' {} {}'.format(user, time))
            training_file.write('\n')
            training_file.flush()

        else:
            testing_file.write('{}'.format(cascade[0][0]))
            for user, time in cascade[1:]:
                if user not in keep_users:
                    continue
                testing_file.write(' {} {}'.format(user, time))
            testing_file.write('\n')
            testing_file.flush()

    training_file.close()
    testing_file.close()
    return keep_users



if __name__ == '__main__':
    keep_users = load_diffusion_data('F:/data/diffusion/digg', 'digg_votes1.csv',)
    load_graph('F:/data/diffusion/digg', 'digg_friends.csv', keep_users)
